# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: xguo

import ujson as json
import math

from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.bitflyer_v1.book_builder import BookBuilder
from coin.base.book.multiple_book_builder_wrapper \
    import MultipleBookBuilderWrapper


# Return None is no valid book date can be found.
def iter_book_data(raw_message):
  payload_message = raw_message.get('M')
  if not payload_message:
    return

  for msg_elem in payload_message:
    assert msg_elem['H'] == 'BFEXHub', msg_elem
    if msg_elem['M'] != 'ReceiveTickers':
      continue
    assert len(msg_elem['A']) == 1, len(msg_elem['A'])
    for book_msg in msg_elem['A'][0]:
      yield book_msg


class BitflyerMultipleBookHandler(WebsocketHandler):
  def __init__(self, aggregated_ticksizes=None):
    super().__init__()
    self._books = {}
    self._aggregated_ticksizes = aggregated_ticksizes

  def on_message(self, record, queue_data, topic_data):
    raw_message = record.data
    if not isinstance(raw_message, dict):
      raw_message = json.loads(raw_message.tobytes())

    native_symbol = None
    book_builder = None
    for message in iter_book_data(raw_message):
      native_symbol = message['pc']  # product code
      if not self.is_subscribed("book", native_symbol):
        continue

      # TODO(xguo): Not sure what is ['tk']['ts'], but these types seems
      # to cause problems.
      if message['tk']['ts'] in ('PREOPEN', 'CLOSED', 'SHUTDOWN'):
        continue
      assert message['tk']['ts'] == 'OPEN'
      if message['tt'] == 'UPDATE':
        book_builder = self._update_book(record.timestamp, native_symbol, message, False)
      elif message['tt'] == 'SNAPSHOT':
        book_builder = self._update_book(record.timestamp, native_symbol, message, True)
      else:
        raise ValueError('Unknown ticker type! %s' % message)

    if native_symbol and book_builder:
      self.publish_event('book', native_symbol, book_builder)

  def _update_book(self, timestamp, native_symbol, message, snapshot):
    if not snapshot:
      if native_symbol not in self._books:
        return
      book_builder = self._books[native_symbol]
      diffs = book_builder.get_original_book().get_diffs(message['bd']['b'], message['bd']['a'])
      for diff in diffs:
        book_builder.update_aggregated_books(diff[0], diff[1], diff[2])
    else:
      book_builder = MultipleBookBuilderWrapper(BookBuilder(native_symbol),
                                                self._aggregated_ticksizes)
      self._books[native_symbol] = book_builder

    book_builder.get_original_book().update(message['bd']['b'], message['bd']['a'])
    book_builder.update_timestamp(timestamp)

    ticker_bid0 = float(message['tk']['bb'])
    ticker_ask0 = float(message['tk']['ba'])

    if not math.isclose(book_builder.bid0().price, ticker_bid0):
      book_builder.get_original_book()._delete_bid_up_to(ticker_bid0)
      book_builder.clear()
      book_builder.generate_aggregated_books()
    if not math.isclose(book_builder.ask0().price, ticker_ask0):
      book_builder.get_original_book()._delete_ask_up_to(ticker_ask0)
      book_builder.clear()
      book_builder.generate_aggregated_books()

    book_ask0 = book_builder.ask0().price
    book_bid0 = book_builder.bid0().price

    assert (book_bid0 < book_ask0)
    assert (book_bid0 <= ticker_bid0)
    assert (book_ask0 >= ticker_ask0)

    if snapshot:
      book_builder.generate_aggregated_books()

    return book_builder
